{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Count Valid Paths in a Tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #depth-first-search #math #dynamic-programming #number-theory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #深度优先搜索 #数学 #动态规划 #数论"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: countPaths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #统计树中的合法路径数目"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一棵 <code>n</code>&nbsp;个节点的无向树，节点编号为&nbsp;<code>1</code>&nbsp;到&nbsp;<code>n</code>&nbsp;。给你一个整数&nbsp;<code>n</code>&nbsp;和一个长度为 <code>n - 1</code>&nbsp;的二维整数数组&nbsp;<code>edges</code>&nbsp;，其中&nbsp;<code>edges[i] = [u<sub>i</sub>, v<sub>i</sub>]</code>&nbsp;表示节点&nbsp;<code>u<sub>i</sub></code> 和&nbsp;<code>v<sub>i</sub></code>&nbsp;在树中有一条边。</p>\n",
    "\n",
    "<p>请你返回树中的 <strong>合法路径数目</strong>&nbsp;。</p>\n",
    "\n",
    "<p>如果在节点 <code>a</code>&nbsp;到节点 <code>b</code>&nbsp;之间 <strong>恰好有一个</strong>&nbsp;节点的编号是质数，那么我们称路径&nbsp;<code>(a, b)</code>&nbsp;是 <strong>合法的</strong>&nbsp;。</p>\n",
    "\n",
    "<p><strong>注意：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li>路径&nbsp;<code>(a, b)</code>&nbsp;指的是一条从节点 <code>a</code>&nbsp;开始到节点 <code>b</code>&nbsp;结束的一个节点序列，序列中的节点 <strong>互不相同</strong>&nbsp;，且相邻节点之间在树上有一条边。</li>\n",
    "\t<li>路径&nbsp;<code>(a, b)</code>&nbsp;和路径&nbsp;<code>(b, a)</code>&nbsp;视为 <strong>同一条</strong>&nbsp;路径，且只计入答案 <strong>一次</strong>&nbsp;。</li>\n",
    "</ul>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong class=\"example\">示例 1：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2023/08/27/example1.png\" style=\"width: 440px; height: 357px;\" /></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]\n",
    "<b>输出：</b>4\n",
    "<b>解释：</b>恰好有一个质数编号的节点路径有：\n",
    "- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。\n",
    "- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。\n",
    "- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。\n",
    "- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。\n",
    "只有 4 条合法路径。\n",
    "</pre>\n",
    "\n",
    "<p><strong class=\"example\">示例 2：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode.com/uploads/2023/08/27/example2.png\" style=\"width: 488px; height: 384px;\" /></p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]\n",
    "<b>输出：</b>6\n",
    "<b>解释：</b>恰好有一个质数编号的节点路径有：\n",
    "- (1, 2) 因为路径 1 到 2 只包含一个质数 2 。\n",
    "- (1, 3) 因为路径 1 到 3 只包含一个质数 3 。\n",
    "- (1, 4) 因为路径 1 到 4 只包含一个质数 2 。\n",
    "- (1, 6) 因为路径 1 到 6 只包含一个质数 3 。\n",
    "- (2, 4) 因为路径 2 到 4 只包含一个质数 2 。\n",
    "- (3, 6) 因为路径 3 到 6 只包含一个质数 3 。\n",
    "只有 6 条合法路径。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>edges.length == n - 1</code></li>\n",
    "\t<li><code>edges[i].length == 2</code></li>\n",
    "\t<li><code>1 &lt;= u<sub>i</sub>, v<sub>i</sub> &lt;= n</code></li>\n",
    "\t<li>输入保证&nbsp;<code>edges</code>&nbsp;形成一棵合法的树。</li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [count-valid-paths-in-a-tree](https://leetcode.cn/problems/count-valid-paths-in-a-tree/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [count-valid-paths-in-a-tree](https://leetcode.cn/problems/count-valid-paths-in-a-tree/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['5\\n[[1,2],[1,3],[2,4],[2,5]]', '6\\n[[1,2],[1,3],[2,4],[3,5],[3,6]]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class DSU:\n",
    "    def __init__(self, n):\n",
    "        self.loc = list(range(n))\n",
    "        self.size = [1] * n\n",
    "\n",
    "    def find(self, x):\n",
    "        if self.loc[x] != x:\n",
    "            self.loc[x] = self.find(self.loc[x])\n",
    "        return self.loc[x]\n",
    "\n",
    "    def union(self, x, y):\n",
    "        a, b = self.find(x), self.find(y)\n",
    "        if a != b:\n",
    "            if self.size[a] >= self.size[b]:\n",
    "                self.loc[b] = a\n",
    "                self.size[a] += self.size[b]\n",
    "            else:\n",
    "                self.loc[a] = b\n",
    "                self.size[b] += self.size[a]\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        \n",
    "        def is_prime(x):\n",
    "            if x == 2:\n",
    "                return True\n",
    "            if not x % 2:\n",
    "                return False\n",
    "            \n",
    "            for j in range(3, int(x**0.5)+1, 2):\n",
    "                if x % j == 0:\n",
    "                    return False\n",
    "            return True\n",
    "                \n",
    "        primes = set(x for x in range(2, n+1) if is_prime(x))\n",
    "        \n",
    "        \n",
    "        dsu = DSU(n+1)\n",
    "        \n",
    "        for u, v in edges:\n",
    "            if u not in primes and v not in primes:\n",
    "                dsu.union(u, v)\n",
    "        \n",
    "        cnt = [1] * (n + 1)\n",
    "        ans = 0\n",
    "        for u, v in edges:\n",
    "            if int(u in primes) + int(v in primes) == 1:\n",
    "                if v in primes:\n",
    "                    u, v = v, u\n",
    "                ans += cnt[u] * dsu.size[dsu.find(v)]\n",
    "                cnt[u] += dsu.size[dsu.find(v)]        \n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        def mul(x, y):\n",
    "            return x * y\n",
    "        \n",
    "        def dfs(x, f, con, prime, r):\n",
    "            v = [1 - prime[x], prime[x]]\n",
    "            for y in con[x]:\n",
    "                if y == f:\n",
    "                    continue\n",
    "                p = dfs(y, x, con, prime, r)\n",
    "                r[0] += mul(p[0], v[1]) + mul(p[1], v[0])\n",
    "                if prime[x]:\n",
    "                    v[1] += p[0]\n",
    "                else:\n",
    "                    v[0] += p[0]\n",
    "                    v[1] += p[1]\n",
    "            return v\n",
    "        \n",
    "        prime = [True] * (n + 1)\n",
    "        prime[1] = False\n",
    "        \n",
    "        all_primes = []\n",
    "        for i in range(2, n + 1):\n",
    "            if prime[i]:\n",
    "                all_primes.append(i)\n",
    "            for x in all_primes:\n",
    "                temp = i * x\n",
    "                if temp > n:\n",
    "                    break\n",
    "                prime[temp] = False\n",
    "        \n",
    "        con = [[] for _ in range(n + 1)]\n",
    "        for e in edges:\n",
    "            con[e[0]].append(e[1])\n",
    "            con[e[1]].append(e[0])\n",
    "        \n",
    "        r = [0]\n",
    "        dfs(1, 0, con, prime, r)\n",
    "        return r[0]\n",
    "            \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        \n",
    "\n",
    "        \n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "MX=10**5+1\n",
    "is_Prime=[True]*MX\n",
    "is_Prime[1]=False\n",
    "for i in range(2,MX):\n",
    "    if is_Prime[i]:\n",
    "        for j in range(i*2,MX,i):\n",
    "            is_Prime[j]=False\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g=[[] for _ in range(n+1)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        ans=0 #方案数\n",
    "        \n",
    "        fa=[i for i in range(n+1)]\n",
    "        def find(x:int)->int:\n",
    "            #fa数组仅在这里出现\n",
    "            if fa[x]!=x:\n",
    "                fa[x]=find(fa[x])\n",
    "            return fa[x]\n",
    "        #并查集做法，预处理？出所有的连通块\n",
    "        route=[int(not is_Prime[i]) for i in range(n+1)]\n",
    "        # def dfs(x:int,pa:int)->None:\n",
    "        #     for y in g[x]:\n",
    "        #         if y!=pa:\n",
    "        #             dfs(y,x)\n",
    "        for x,y in edges:\n",
    "            if not is_Prime[x] and not is_Prime[y]:\n",
    "                fx=find(x)\n",
    "                fy=find(y)\n",
    "                # if fx==fy:continue\n",
    "                fa[fy]=fx\n",
    "                route[fx]+=route[fy]\n",
    "        # dfs(1,-1)\n",
    "        for x in range(1,n+1):\n",
    "            if not is_Prime[x]:continue\n",
    "            res=1\n",
    "            for y in g[x]:\n",
    "                if not is_Prime[y]:\n",
    "                    ans+=res*route[find(y)]\n",
    "                    res+=route[find(y)]\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "MX=10**5+1\n",
    "is_Prime=[True]*MX\n",
    "is_Prime[1]=False\n",
    "for i in range(2,MX):\n",
    "    if is_Prime[i]:\n",
    "        for j in range(i*2,MX,i):\n",
    "            is_Prime[j]=False\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g=[[] for _ in range(n+1)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        ans=0 #方案数\n",
    "        \n",
    "        fa=[i for i in range(n+1)]\n",
    "        def find(x:int)->int:\n",
    "            if fa[x]!=x:\n",
    "                fa[x]=find(fa[x])\n",
    "            return fa[x]\n",
    "        #并查集做法，预处理？出所有的连通块\n",
    "        route=[int(not is_Prime[i]) for i in range(n+1)]\n",
    "        # def dfs(x:int,pa:int)->None:\n",
    "        #     for y in g[x]:\n",
    "        #         if y!=pa:\n",
    "        #             dfs(y,x)\n",
    "        for x,y in edges:\n",
    "            if not is_Prime[x] and not is_Prime[y]:\n",
    "                fx=find(x)\n",
    "                fy=find(y)\n",
    "                if fx==fy:continue\n",
    "                fa[fy]=fx\n",
    "                route[fx]+=route[fy]\n",
    "        # dfs(1,-1)\n",
    "        for x in range(1,n+1):\n",
    "            if not is_Prime[x]:continue\n",
    "            res=1\n",
    "            for y in g[x]:\n",
    "                if not is_Prime[y]:\n",
    "                    ans+=res*route[find(y)]\n",
    "                    res+=route[find(y)]\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "import math\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def prime(self, n):\n",
    "        is_prime = [True] * n\n",
    "        is_prime[1] = False\n",
    "        for i in range(2, int(math.sqrt(n)) + 1):\n",
    "            if is_prime[i]:\n",
    "                for j in range(i * i, n, i):\n",
    "                    is_prime[j] = False\n",
    "\n",
    "        return is_prime\n",
    "\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        if not edges:\n",
    "            return 0\n",
    "        max_num = max([max(item) for item in edges])\n",
    "        is_prime: [bool] = self.prime(max_num + 1)\n",
    "        g = [[] for _ in range(n + 1)]\n",
    "        for edge in edges:\n",
    "            g[edge[0]].append(edge[1])\n",
    "            g[edge[1]].append(edge[0])\n",
    "\n",
    "        def dfs(nod, val):\n",
    "            nodes.append(nod)\n",
    "            for nod_e in g[nod]:\n",
    "                if is_prime[nod_e] or nod_e == val:\n",
    "                    continue\n",
    "                dfs(nod_e, nod)\n",
    "\n",
    "        ans = 0\n",
    "        size = [0] * (n + 1)  # 节点联通的非质数点的数量\n",
    "        for edge in range(1, n + 1):\n",
    "            if not is_prime[edge]:\n",
    "                continue\n",
    "            edge_size = 0  # 质数edge联通的块中的所有的非质数的数量\n",
    "            for e in g[edge]:  # 循环联通块\n",
    "                if is_prime[e]:\n",
    "                    continue\n",
    "\n",
    "                if size[e] == 0:  # 统计节点联通的非质数的数量\n",
    "                    nodes = []\n",
    "                    dfs(e, -1)\n",
    "                    length = len(nodes)\n",
    "                    for node in nodes:\n",
    "                        size[node] = length\n",
    "\n",
    "                ans += size[e] * edge_size  # [edge1、edge2] - e - [edge3、edge4] => 以e为线路中的节点，数量为 2 * 2\n",
    "                edge_size += size[e]\n",
    "\n",
    "            ans += edge_size  # 最后加上以e为端点的线路\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "MAXN = 100001\n",
    "\n",
    "prims = []\n",
    "is_prim = [True] * MAXN\n",
    "min_prim = [None] * MAXN\n",
    "\n",
    "is_prim[1] = False\n",
    "\n",
    "for i in range(2, MAXN):\n",
    "    if is_prim[i]:\n",
    "        prims.append(i)\n",
    "        min_prim[i] = i\n",
    "    for p in prims:\n",
    "        if p > min_prim[i] or p * i >= MAXN:\n",
    "            break\n",
    "        v = p * i\n",
    "        is_prim[v] = False\n",
    "        min_prim[v] = p\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        graph = [[] for _ in range(n + 1)]\n",
    "        for a, b in edges:\n",
    "            graph[a].append(b)\n",
    "            graph[b].append(a)\n",
    "        ans = 0\n",
    "        def dfs(x, p):\n",
    "            nonlocal ans\n",
    "            c0 = 0\n",
    "            c1 = 0\n",
    "            prim = is_prim[x]\n",
    "            if prim:\n",
    "                c1 += 1\n",
    "            else:\n",
    "                c0 += 1\n",
    "            for b in graph[x]:\n",
    "                if b != p:\n",
    "                    cc0, cc1 = dfs(b, x)\n",
    "                    if prim:\n",
    "                        ans += cc0 * c1\n",
    "                        c1 += cc0\n",
    "                    else:\n",
    "                        ans += cc1 * c0\n",
    "                        ans += cc0 * c1\n",
    "                        c0 += cc0\n",
    "                        c1 += cc1\n",
    "            return c0, c1\n",
    "        dfs(1, None)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "       \n",
    "        gsize = [0]*(n+1)      \n",
    "       \n",
    "        prime = set()\n",
    "        nums = [0] * (n+1)\n",
    "        nums[1] = 1\n",
    "        for i in range(2, n+1):\n",
    "            if nums[i] == 0:\n",
    "                prime.add(i)\n",
    "                p = i\n",
    "                m = n // p\n",
    "                for j in range(1, m + 1):\n",
    "                    nums[j * p] = 1\n",
    "       \n",
    "        \n",
    "        adj=defaultdict(list)\n",
    "        for a,b in edges:\n",
    "            adj[a].append(b)\n",
    "            adj[b].append(a)\n",
    "\n",
    "        def groupsize(p:int)->int:\n",
    "            if gsize[p]>0:\n",
    "                return gsize[p]\n",
    "            visit=set([p])\n",
    "            dfs=[p]\n",
    "            while dfs:\n",
    "                cur=dfs.pop()\n",
    "                for d in adj[cur]:\n",
    "                    if d not in visit and d not in prime:\n",
    "                        visit.add(d)\n",
    "                        dfs.append(d)\n",
    "            cnt=len(visit)\n",
    "            for v in visit:\n",
    "                gsize[v]=cnt\n",
    "            return cnt\n",
    "        ans=0\n",
    "        for p in prime:\n",
    "            conp=adj[p]\n",
    "            cnts=[]\n",
    "            for cp in conp:\n",
    "                if cp not in prime:          \n",
    "                    cnts.append(groupsize(cp))\n",
    "            if len(cnts)>0:             \n",
    "                presum=cnts[0]\n",
    "                for i in range(1,len(cnts)):\n",
    "                    ans+=cnts[i]*presum\n",
    "                    presum+=cnts[i]\n",
    "                ans+=presum\n",
    "        return ans\n",
    "          \n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "MX=10**5+1\n",
    "is_Prime=[True]*MX\n",
    "is_Prime[1]=False\n",
    "for i in range(2,MX):\n",
    "    if is_Prime[i]:\n",
    "        for j in range(i*2,MX,i):\n",
    "            is_Prime[j]=False\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g=[[] for _ in range(n+1)]\n",
    "        for x,y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        ans=0 #方案数\n",
    "        \n",
    "        fa=[i for i in range(n+1)]\n",
    "        def find(x:int)->int:\n",
    "            #fa数组仅在这里出现\n",
    "            if fa[x]!=x:\n",
    "                fa[x]=find(fa[x])\n",
    "            return fa[x]\n",
    "        #并查集做法，预处理？出所有的连通块\n",
    "        route=[int(not is_Prime[i]) for i in range(n+1)]\n",
    "        def dfs(x:int,pa:int)->None:\n",
    "            for y in g[x]:\n",
    "                if y!=pa:\n",
    "                    dfs(y,x)\n",
    "                    if not is_Prime[x] and not is_Prime[y]:\n",
    "                        fx=find(x)\n",
    "                        fy=find(y)\n",
    "                        # if fx==fy:continue\n",
    "                        fa[fy]=fx\n",
    "                        route[fx]+=route[fy]\n",
    "\n",
    "        '''\n",
    "        #每次都一定有一个新的点到来，不可能两个点都是之前遇到过的\n",
    "        for x,y in edges:\n",
    "            if not is_Prime[x] and not is_Prime[y]:\n",
    "                fx=find(x)\n",
    "                fy=find(y)\n",
    "                # if fx==fy:continue\n",
    "                fa[fy]=fx\n",
    "                route[fx]+=route[fy]\n",
    "        '''\n",
    "        dfs(1,-1)\n",
    "        for x in range(1,n+1):\n",
    "            if not is_Prime[x]:continue\n",
    "            res=1\n",
    "            for y in g[x]:\n",
    "                if not is_Prime[y]:\n",
    "                    ans+=res*route[find(y)]\n",
    "                    res+=route[find(y)]\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "N = 100010\n",
    "primes = [0]*(N)\n",
    "st = [0]*(N)\n",
    "cnt = 0\n",
    "st[1] = 1\n",
    "\n",
    "for i in range(2,N):\n",
    "    if not st[i]:\n",
    "        primes[cnt] = i\n",
    "        cnt+=1\n",
    "    for j in range(N):\n",
    "        if primes[j]*i>=N or primes[j]*i==0:break\n",
    "        st[primes[j]*i] = 1\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "\n",
    "        def dfs(u,fa):\n",
    "            node.append(u)\n",
    "            for v in g[u]:\n",
    "                if v==fa:continue\n",
    "                if st[v]==0:continue\n",
    "                dfs(v,u)\n",
    "\n",
    "\n",
    "        g = defaultdict(list)\n",
    "        \n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        sz = [0]*(n+10)\n",
    "        res = 0\n",
    "        \n",
    "        for i in range(1,n+1):\n",
    "            if st[i]:continue\n",
    "\n",
    "            sums = 0\n",
    "            for v in g[i]:\n",
    "                if not st[v]:continue\n",
    "                if not sz[v]:\n",
    "                    node = []\n",
    "                    dfs(v,-1)\n",
    "                    for z in node:\n",
    "                        sz[z] = len(node)\n",
    "                res+=sums*sz[v]\n",
    "                sums+=sz[v]\n",
    "            res+=sums\n",
    "        \n",
    "        return res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "\n",
    "        # non bineary\n",
    "        def check_prime(v):\n",
    "\n",
    "            if v == 1:\n",
    "                return False\n",
    "\n",
    "            for i in range(2, int(sqrt(v) + 1)):\n",
    "                if v % i == 0:\n",
    "                    #print(f'check prime {v} False')\n",
    "                    return False\n",
    "\n",
    "            #print(f'check prime {v} True')\n",
    "            return True\n",
    "\n",
    "        # ret = 0\n",
    "        # def dfs(n):\n",
    "\n",
    "        #     if not n:\n",
    "        #         return (0, 0)\n",
    "            \n",
    "        #     l_all, l_valid = dfs(n.left)\n",
    "        #     r_all, r_valid = dfs(n.right)\n",
    "\n",
    "        #     all_ret = 0\n",
    "        #     valid_ret = 0\n",
    "        #     if check_prime(n.val):\n",
    "        #         valid = (l_all + 1) * (r_all + 1)\n",
    "        #         valid_ret += l_all + r_all + 1\n",
    "        #     else:\n",
    "        #         valid = (l_all + 1) * (l_ret + 1) + (r_all + 1) * (l_ret + 1)\n",
    "        #         valid_ret += l_ret + r_ret\n",
    "        #         all_ret += l_all + r_all + 1\n",
    "            \n",
    "        #     nonlocal ret\n",
    "        #     ret += valid\n",
    "\n",
    "        #     return all_ret, valid_ret\n",
    "        \n",
    "            \n",
    "        d = defaultdict(list)\n",
    "        for s, e in edges:\n",
    "            d[s].append(e)\n",
    "            d[e].append(s)\n",
    "        \n",
    "        # cur = [1]\n",
    "        # while cur:\n",
    "        #     next = []\n",
    "        #     for p in cur:\n",
    "        #         for c in d[p]:\n",
    "        #             d[c].remove(p)\n",
    "\n",
    "        ret = 0\n",
    "        visited = set()\n",
    "        def dfs(n):\n",
    "\n",
    "            if n in visited:\n",
    "                return 0, 0\n",
    "            \n",
    "            visited.add(n)\n",
    "\n",
    "            nonlocal ret\n",
    "\n",
    "            all_common = 0\n",
    "            all_valid = 0\n",
    "\n",
    "            pre_valid = 0\n",
    "            pre_common = 0\n",
    "            if check_prime(n):\n",
    "                all_valid += 1\n",
    "                for c in d[n]:\n",
    "                    valid, common = dfs(c)\n",
    "                    ret += common\n",
    "                    ret += pre_common * common\n",
    "\n",
    "                    pre_valid += valid\n",
    "                    pre_common += common\n",
    "                    \n",
    "                    all_valid += common\n",
    "            else:\n",
    "                all_common += 1\n",
    "                for c in d[n]:\n",
    "                    valid, common = dfs(c)\n",
    "                    ret += valid\n",
    "                    ret += pre_common * valid + pre_valid * common\n",
    "\n",
    "                    pre_valid += valid\n",
    "                    pre_common += common\n",
    "\n",
    "                    all_valid += valid\n",
    "                    all_common += common\n",
    "            \n",
    "            #print(f'{n} {all_valid} {all_common}')\n",
    "            return all_valid, all_common\n",
    "        \n",
    "        # nodes = set(range(1, n + 1))\n",
    "        # for i, arr in d.items():\n",
    "        #     print(f'{nodes} {arr}')\n",
    "        #     nodes = nodes - set(arr)\n",
    "        #     print(f'{nodes} {arr}')\n",
    "\n",
    "        # assert len(nodes) == 1,  nodes\n",
    "        # dfs(nodes.pop())\n",
    "\n",
    "        dfs(1)\n",
    "\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        if n < 2:\n",
    "            return 0 \n",
    "        isPrime = [1] * (n+1) \n",
    "        isPrime[0] = isPrime[1] = 0  \n",
    "        for i in range(2, int(math.sqrt(n)) + 2):\n",
    "            if isPrime[i]:\n",
    "                for j in range(i*2, n+1, i):\n",
    "                    isPrime[j] = 0 \n",
    "        \n",
    "        g = defaultdict(list) \n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a) \n",
    "\n",
    "        def dfs(x, fa, node):\n",
    "            node.append(x) \n",
    "            tot = 1 \n",
    "            for y in g[x]:\n",
    "                if y != fa and not isPrime[y]:\n",
    "                    tot += dfs(y, x, node)  \n",
    "            return tot \n",
    "            \n",
    "        dp = defaultdict(int) \n",
    "        def func(root):\n",
    "            arr = []\n",
    "            tot = 0 \n",
    "            for x in g[root]:\n",
    "                if not isPrime[x]:\n",
    "                    if x in dp:\n",
    "                        arr.append(dp[x]) \n",
    "                        continue \n",
    "                    node = []\n",
    "                    res = dfs(x, root, node)\n",
    "                    for y in node:\n",
    "                        dp[y] = res \n",
    "                    arr.append(res) \n",
    "            m = len(arr) \n",
    "            pre = [0] * (m+1)\n",
    "            for i in range(1, m+1):\n",
    "                pre[i] = pre[i-1] + arr[i-1] \n",
    "            for i in range(1, m+1):\n",
    "                tot += pre[i-1] * arr[i-1] + arr[i-1] \n",
    "            return tot \n",
    "        \n",
    "        ans = 0 \n",
    "        for i in range(2, n+1):\n",
    "            if isPrime[i]:\n",
    "                ans += func(i) \n",
    "        return ans "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for u,v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        def get_p(n):\n",
    "            primes= []\n",
    "            st = [False] * (n+1)\n",
    "            for i in range(2, n+1):\n",
    "                if not st[i]:\n",
    "                    primes.append(i)\n",
    "                for p in primes:\n",
    "                    if p* i > n: break\n",
    "                    st[p*i]= True\n",
    "                    if i%p==0:\n",
    "                        break\n",
    "            return primes\n",
    "        primes = set(get_p(n))\n",
    "        def dfs(u):\n",
    "            vis.add(u)\n",
    "            cc.append(u)\n",
    "            for v in g[u]:\n",
    "                if v not in vis and v not in primes:\n",
    "                    dfs(v)\n",
    "        vis = set()\n",
    "        cc_cntr = [0] * (n+1)\n",
    "        for u in range(1,n+1):\n",
    "            if u not in primes and u not in vis:\n",
    "                cc = []\n",
    "                dfs(u)\n",
    "                for v in cc:\n",
    "                    cc_cntr[v] = len(cc)\n",
    "\n",
    "        ans = 0\n",
    "        for p in primes:\n",
    "            rs = 0\n",
    "            rrs = 0\n",
    "            for v in g[p]:\n",
    "                if v not in primes:\n",
    "                    v = cc_cntr[v]\n",
    "                    rs += v\n",
    "                    rrs += v*v\n",
    "            ans += rs + (rs*rs - rrs)//2\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        #start = time.time()\n",
    "        \n",
    "        # primary\n",
    "        prim = []\n",
    "        isprim = [False] * (n+1)\n",
    "        for i in range(2, n+1):\n",
    "            for p in prim:\n",
    "                if i % p == 0:\n",
    "                    break\n",
    "            else:\n",
    "                prim.append(i)\n",
    "                isprim[i] = True\n",
    "        #print(time.time()-start, 'prim', prim)\n",
    "        \n",
    "        # parent\n",
    "        parent = [i for i in range(n+1)]\n",
    "        def find(x):\n",
    "            if parent[x] == x:\n",
    "                return x\n",
    "            parent[x] = find(parent[x])\n",
    "            #print('find', x, parent[x])\n",
    "            return parent[x]\n",
    "        \n",
    "        def union(a, b):\n",
    "            fa, fb = find(a), find(b)\n",
    "            #print('union', a, b,'=', fa, fb)\n",
    "            if fa == fb:\n",
    "                return\n",
    "            if fa < fb:\n",
    "                parent[fb] = fa\n",
    "            else:\n",
    "                parent[fa] = fb\n",
    "        \n",
    "        for a, b in edges:\n",
    "            if not isprim[a] and not isprim[b]:\n",
    "                union(a, b)\n",
    "        #print(time.time()-start, 'parent union', parent)\n",
    "\n",
    "        # count not prime sets\n",
    "        setcount = [0] * (n+1)\n",
    "        for i in range(1, n+1):\n",
    "            if not isprim[i]:\n",
    "                setcount[find(i)] += 1\n",
    "        #print(time.time()-start, 'count prime sets', setcount)\n",
    "        \n",
    "        # count linked\n",
    "        linked = [set() for _ in range(n+1)]\n",
    "        for a, b in edges:\n",
    "            aa, bb = isprim[a], isprim[b]\n",
    "            #print('link', a,b,'=',aa,bb)\n",
    "            if aa == bb:\n",
    "                continue\n",
    "            if aa:\n",
    "                linked[a].add(parent[b])\n",
    "            else:\n",
    "                linked[b].add(parent[a])\n",
    "        #print(time.time()-start, 'count linked', linked)\n",
    "        \n",
    "        # count final\n",
    "        ssum = 0\n",
    "        for i in prim:\n",
    "            vcount = [setcount[v] for v in linked[i]]\n",
    "            #print('count', i, vcount)\n",
    "            allsum = 0\n",
    "            for v in vcount:\n",
    "                ssum += v * allsum\n",
    "                allsum += v\n",
    "            ssum += allsum\n",
    "\n",
    "        return ssum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        s = [0] * (n + 1)\n",
    "        s[1] = 1\n",
    "        for i in range(2, n + 1):\n",
    "            if s[i] == 0:\n",
    "                for j in range(i * i, n + 1, i):\n",
    "                    s[j] = 1\n",
    "        e = collections.defaultdict(list)\n",
    "        for x, y in edges:\n",
    "            e[x].append(y)\n",
    "            e[y].append(x)\n",
    "        h = set()\n",
    "        f = {}\n",
    "        self.ret = 0\n",
    "\n",
    "        def dfs(x):\n",
    "            h.add(x)\n",
    "            a = []\n",
    "            for y in e[x]:\n",
    "                if y in h:\n",
    "                    continue\n",
    "                a.append(dfs(y))\n",
    "            if s[x] == 0:\n",
    "                v = 0\n",
    "                for z in a:\n",
    "                    self.ret += v * z\n",
    "                    v += z\n",
    "                self.ret += v\n",
    "                f[x] = 0\n",
    "            else:\n",
    "                f[x] = sum(a) + 1\n",
    "            return f[x]\n",
    "\n",
    "        dfs(1)\n",
    "\n",
    "        h = set()\n",
    "\n",
    "        def dfs2(x, c):\n",
    "            h.add(x)\n",
    "            if s[x] == 0:\n",
    "                self.ret += c\n",
    "            now = 0\n",
    "            for y in e[x]:\n",
    "                if y in h:\n",
    "                    continue\n",
    "                if s[x] == 0:\n",
    "                    self.ret += c * f[y]\n",
    "                else:\n",
    "                    now += f[y]\n",
    "            for y in e[x]:\n",
    "                if y in h:\n",
    "                    continue\n",
    "                if s[x] == 1:\n",
    "                    dfs2(y, c + 1 + now - f[y])\n",
    "                else:\n",
    "                    dfs2(y, 0)\n",
    "\n",
    "        dfs2(1, 0)\n",
    "        return self.ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Primer(object):\n",
    "    \"\"\"素数相关运算的辅助类\n",
    "\n",
    "    Attributes:\n",
    "        1. n：用素数筛选法提前计算好的最大数\n",
    "        2. primes：n以内的所有素数列表\n",
    "        3. min_prime_factors：每个下标数的最小素因数（最小素因数不存在时为-1）\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    @property\n",
    "    def n(self):\n",
    "        return self._n\n",
    "\n",
    "    @property\n",
    "    def primes(self):\n",
    "        return self._primes\n",
    "\n",
    "    @property\n",
    "    def min_prime_factors(self):\n",
    "        return self._min_prime_factors\n",
    "\n",
    "    def __init__(self, n=100003):\n",
    "        \"\"\"初始化\n",
    "\n",
    "        生成n以内的所有素数primes，同时得到每个数的最小素因数min_prime_factors\n",
    "\n",
    "        Time Complexity:\n",
    "            O(n)\n",
    "\n",
    "        \"\"\"\n",
    "        primes = []\n",
    "        min_prime_factors = [-1, -1] + [0] * (n - 1)\n",
    "        for i in range(2, n + 1):\n",
    "            if min_prime_factors[i] == 0:\n",
    "                primes.append(i)\n",
    "                min_prime_factors[i] = i\n",
    "            j = 0\n",
    "            while i * primes[j] <= n:\n",
    "                min_prime_factors[i * primes[j]] = primes[j]\n",
    "                if i % primes[j] == 0:\n",
    "                    break\n",
    "                j += 1\n",
    "        self._n = n\n",
    "        self._primes = primes\n",
    "        self._min_prime_factors = min_prime_factors\n",
    "        \n",
    "p = Primer()\n",
    "isp = [v == i for i, v in enumerate(p.min_prime_factors)]\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        e = [[] for _ in range(n + 1)]\n",
    "        for a, b in edges:\n",
    "            e[a].append(b)\n",
    "            e[b].append(a)\n",
    "        \n",
    "        @functools.cache\n",
    "        def dfs(t, l):\n",
    "            s = 1\n",
    "            for q in e[t]:\n",
    "                if q != l and not isp[q]:\n",
    "                    s += dfs(q, t)\n",
    "            return s\n",
    "        \n",
    "        def get(lt):\n",
    "            s = sum(lt)\n",
    "            return (s * s - sum([v * v for v in lt])) // 2 + s\n",
    "        \n",
    "        ans = 0\n",
    "        for i in range(1, n + 1):\n",
    "            if isp[i]:\n",
    "                lt = []\n",
    "                for q in e[i]:\n",
    "                    if not isp[q]:\n",
    "                        lt.append(dfs(q, -1))\n",
    "                ans += get(lt)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        flags = [1 for _ in range(n+1)]   \n",
    "        ## 2,3,5, is prime    1-prime, -1 not prime\n",
    "        flags[0] = 0\n",
    "        primes = []\n",
    "        for i in range(2,n+1):\n",
    "            if flags[i]==1:\n",
    "                primes.append(i)\n",
    "            for pv in primes:\n",
    "                if pv*i>n: break\n",
    "                flags[pv*i] = 0\n",
    "                if i%pv==0:\n",
    "                    break\n",
    "        \n",
    "        # construct adjacent relation\n",
    "        import collections\n",
    "        rel = collections.defaultdict(list)\n",
    "        for s,e in edges:\n",
    "            rel[s].append(e)\n",
    "            rel[e].append(s)\n",
    "        # print(len(primes))\n",
    "        ## primes \n",
    "        res = 0\n",
    "        primes2 = set(primes)\n",
    "        # print(primes)\n",
    "        saved_dfs = {}\n",
    "        def countpath(node):           \n",
    "            def dfs(dnode, farther):\n",
    "                if dnode in primes2: return 0\n",
    "                if (dnode, farther) in saved_dfs: \n",
    "                    return saved_dfs[(dnode, farther)]\n",
    "                tres = 1\n",
    "                for tnode in rel[dnode]:\n",
    "                    if tnode==farther: continue\n",
    "                    tres += dfs(tnode, dnode)\n",
    "                saved_dfs[(dnode, farther)] = tres\n",
    "                return tres\n",
    "            \n",
    "            sons = []\n",
    "            for snode in rel[node]:\n",
    "                tc = dfs(snode, -1)\n",
    "                if tc!=0:\n",
    "                    sons.append(tc)\n",
    "            if len(sons)==0: return 0\n",
    "            fsum = sum(sons)\n",
    "            eres = 0\n",
    "            for tv in sons:\n",
    "                eres += tv\n",
    "                eres += tv*(fsum-tv)\n",
    "                fsum -= tv\n",
    "            # print(sons, node, eres)\n",
    "            return eres\n",
    "\n",
    "        \n",
    "        for tv in primes:\n",
    "            res += countpath(tv)\n",
    "            print(tv, res)\n",
    "        \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Bitmap:\n",
    "    def __init__(self, n):\n",
    "        self.bm = bytearray((n + 1) // 2)\n",
    "    def set(self, i):\n",
    "        self.bm[i >> 1] |= 1 << (i & 1)\n",
    "    def unset(self, i):\n",
    "        self.bm[i >> 1] &= ~(1 << (i & 1))\n",
    "    def get(self, i):\n",
    "        return self.bm[i >> 1] & (1 << (i & 1))\n",
    "\n",
    "\n",
    "# Reference: https://stackoverflow.com/a/3035188\n",
    "import numpy as np\n",
    "def primesfrom2to(n):\n",
    "    \"\"\"Input n>=6, Returns a array of primes, 2 <= p < n\"\"\"\n",
    "    sieve = np.ones(n // 3 + (n % 6 == 2), dtype=bool)\n",
    "    for i in range(1, int(n**0.5) // 3 + 1):\n",
    "        if sieve[i]:\n",
    "            k = 3 * i + 1 | 1\n",
    "            sieve[k * k // 3 :: 2 * k] = False\n",
    "            sieve[k * (k - 2 * (i & 1) + 4) // 3 :: 2 * k] = False\n",
    "    return np.r_[2, 3, ((3 * np.nonzero(sieve)[0][1:] + 1) | 1)]\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def __init__(self) -> None:\n",
    "        MAXN = 100_010\n",
    "        self.is_prime = Bitmap(MAXN)\n",
    "        primes = primesfrom2to(MAXN)\n",
    "        for prime in primes:\n",
    "            self.is_prime.set(prime)\n",
    "\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        adj_list = {}\n",
    "        for u, v in edges:\n",
    "            if u not in adj_list:\n",
    "                adj_list[u] = []\n",
    "            if v not in adj_list:\n",
    "                adj_list[v] = []\n",
    "            adj_list[u].append(v)\n",
    "            adj_list[v].append(u)\n",
    "\n",
    "        # 树上DP\n",
    "        f = [0] * (n + 1)  # f[i] 表示 以i为根的子树中，所有节点均是合数 的路径数\n",
    "        g = [0] * (n + 1)  # g[i] 表示 以i为根的子树中，恰有一个节点是质数 的路径数\n",
    "        \n",
    "        self.ans = 0\n",
    "\n",
    "        def dfs(cur, pre):\n",
    "            f[cur], g[cur] = 1, 0\n",
    "            if self.is_prime.get(cur):\n",
    "                f[cur], g[cur] = g[cur], f[cur]\n",
    "            if cur in adj_list:\n",
    "                for child in adj_list[cur]:\n",
    "                    if child != pre:  # 防止重复遍历\n",
    "                        dfs(child, cur)\n",
    "                        self.ans += f[cur] * g[child] + g[cur] * f[child]\n",
    "                        if self.is_prime.get(cur):\n",
    "                            g[cur] += f[child]\n",
    "                        else:\n",
    "                            f[cur] += f[child]\n",
    "                            g[cur] += g[child]\n",
    "\n",
    "        dfs(1, 0)\n",
    "        return self.ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List, Tuple\n",
    "from typing import Callable, Generic, List, TypeVar\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "\n",
    "class Rerooting(Generic[T]):\n",
    "    __slots__ = (\"adjList\", \"_n\", \"_decrement\")\n",
    "\n",
    "    def __init__(self, n: int, decrement: int = 0):\n",
    "        self.adjList = [[] for _ in range(n)]\n",
    "        self._n = n\n",
    "        self._decrement = decrement\n",
    "\n",
    "    def addEdge(self, u: int, v: int) -> None:\n",
    "        u -= self._decrement\n",
    "        v -= self._decrement\n",
    "        self.adjList[u].append(v)\n",
    "        self.adjList[v].append(u)\n",
    "\n",
    "    def rerooting(\n",
    "        self,\n",
    "        e: Callable[[int], T],\n",
    "        op: Callable[[T, T], T],\n",
    "        composition: Callable[[T, int, int, int], T],\n",
    "        root=0,\n",
    "    ) -> List[\"T\"]:\n",
    "        root -= self._decrement\n",
    "        assert 0 <= root < self._n\n",
    "        parents = [-1] * self._n\n",
    "        order = [root]\n",
    "        stack = [root]\n",
    "        while stack:\n",
    "            cur = stack.pop()\n",
    "            for next in self.adjList[cur]:\n",
    "                if next == parents[cur]:\n",
    "                    continue\n",
    "                parents[next] = cur\n",
    "                order.append(next)\n",
    "                stack.append(next)\n",
    "\n",
    "        dp1 = [e(i) for i in range(self._n)]\n",
    "        dp2 = [e(i) for i in range(self._n)]\n",
    "        for cur in order[::-1]:\n",
    "            res = e(cur)\n",
    "            for next in self.adjList[cur]:\n",
    "                if parents[cur] == next:\n",
    "                    continue\n",
    "                dp2[next] = res\n",
    "                res = op(res, composition(dp1[next], cur, next, 0))\n",
    "            res = e(cur)\n",
    "            for next in self.adjList[cur][::-1]:\n",
    "                if parents[cur] == next:\n",
    "                    continue\n",
    "                dp2[next] = op(res, dp2[next])\n",
    "                res = op(res, composition(dp1[next], cur, next, 0))\n",
    "            dp1[cur] = res\n",
    "\n",
    "        for newRoot in order[1:]:\n",
    "            parent = parents[newRoot]\n",
    "            dp2[newRoot] = composition(op(dp2[newRoot], dp2[parent]), parent, newRoot, 1)\n",
    "            dp1[newRoot] = op(dp1[newRoot], dp2[newRoot])\n",
    "        return dp1\n",
    "\n",
    "from collections import Counter\n",
    "\n",
    "\n",
    "class EratosthenesSieve:\n",
    "    \"\"\"埃氏筛\"\"\"\n",
    "\n",
    "    __slots__ = \"_minPrime\"  # 每个数的最小质因数\n",
    "\n",
    "    def __init__(self, maxN: int):\n",
    "        \"\"\"预处理 O(nloglogn)\"\"\"\n",
    "        minPrime = list(range(maxN + 1))\n",
    "        upper = int(maxN**0.5) + 1\n",
    "        for i in range(2, upper):\n",
    "            if minPrime[i] < i:\n",
    "                continue\n",
    "            for j in range(i * i, maxN + 1, i):\n",
    "                if minPrime[j] == j:\n",
    "                    minPrime[j] = i\n",
    "        self._minPrime = minPrime\n",
    "\n",
    "    def isPrime(self, n: int) -> bool:\n",
    "        if n < 2:\n",
    "            return False\n",
    "        return self._minPrime[n] == n\n",
    "\n",
    "    def getPrimeFactors(self, n: int) -> \"Counter[int]\":\n",
    "        \"\"\"求n的质因数分解 O(logn)\"\"\"\n",
    "        res, f = Counter(), self._minPrime\n",
    "        while n > 1:\n",
    "            m = f[n]\n",
    "            res[m] += 1\n",
    "            n //= m\n",
    "        return res\n",
    "\n",
    "    def getPrimes(self) -> List[int]:\n",
    "        return [x for i, x in enumerate(self._minPrime) if i >= 2 and i == x]\n",
    "\n",
    "\n",
    "P = EratosthenesSieve(int(1e5) + 10)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        E = Tuple[int, int]\n",
    "\n",
    "        def e(root: int) -> E:\n",
    "            return (0, 0)  # 幺元 (恰好0个质数的路径数, 恰好1个质数的路径数)\n",
    "\n",
    "        def op(childRes1: E, childRes2: E) -> E:\n",
    "            return (childRes1[0] + childRes2[0], childRes1[1] + childRes2[1])\n",
    "\n",
    "        def composition(fromRes: E, parent: int, cur: int, direction: int) -> E:\n",
    "            \"\"\"direction: 0: cur -> parent, 1: parent -> cur\"\"\"\n",
    "            from_ = cur if direction == 0 else parent\n",
    "            to_ = parent if direction == 0 else cur\n",
    "            isPrime = int(P.isPrime(from_ + 1))\n",
    "            zero, one = fromRes\n",
    "            return (0, zero + 1) if isPrime else (zero + 1, one)\n",
    "\n",
    "        R = Rerooting(n)\n",
    "        for u, v in edges:\n",
    "            R.addEdge(u - 1, v - 1)\n",
    "        dp = R.rerooting(e=e, op=op, composition=composition, root=0)  # !答案不包含当前根节点\n",
    "        res = 0\n",
    "        for root in range(n):\n",
    "            if P.isPrime(root + 1):\n",
    "                res += dp[root][0]\n",
    "            else:\n",
    "                res += dp[root][1]\n",
    "        return res // 2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "isP = [1] * (10 ** 5 + 1)\n",
    "isP[0] = isP[1] = 0\n",
    "for i in range(2, 10 ** 5):\n",
    "    if isP[i] == 0: continue\n",
    "    for j in range(i, 10 ** 5):\n",
    "        if i * j > 10 ** 5: break\n",
    "        isP[i * j] = 0\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        gra = defaultdict(set)\n",
    "        for a, b in edges:\n",
    "            gra[a].add(b)\n",
    "            gra[b].add(a)\n",
    "\n",
    "\n",
    "        res = 0\n",
    "        def dfs(node, pre):\n",
    "            nonlocal res\n",
    "            \n",
    "            r = [0, 0]\n",
    "            if isP[node]:\n",
    "                r[1] += 1\n",
    "            else:\n",
    "                r[0] += 1\n",
    "            \n",
    "            for n in gra[node]:\n",
    "                if n == pre: continue\n",
    "                r0, r1 = dfs(n, node)\n",
    "                res += r0 * r[1] + r1 * r[0]\n",
    "                if isP[node]: \n",
    "                    r[1] += r0\n",
    "                else:\n",
    "                    r[0] += r0\n",
    "                    r[1] += r1\n",
    "            return r\n",
    "        dfs(1, -1)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "# 埃氏筛模板\n",
    "def eratosthenes(n):\n",
    "  primes = []\n",
    "  is_prime = [True] * (n + 1)\n",
    "  for i in range(2, n+1):\n",
    "      if is_prime[i]:\n",
    "          primes.append(i)\n",
    "          # 用当前素数i去筛掉所有能被它整除的数\n",
    "          for j in range(i * 2, n+1, i):\n",
    "              is_prime[j] = False\n",
    "  return primes\n",
    "\n",
    "primes = eratosthenes(int(1e5))\n",
    "\n",
    "# 并查集模板\n",
    "class UnionFind:\n",
    "  \n",
    "  def __init__(self, n):\n",
    "      self.arr = list(range(n))\n",
    "      self.wight = [1] * n\n",
    "      \n",
    "  def getRoot(self, n):\n",
    "      m = n\n",
    "      arr = self.arr\n",
    "      while n != arr[n]: \n",
    "          n = arr[n]\n",
    "      arr[m] = n\n",
    "      return n\n",
    "  \n",
    "  def union(self, a, b):\n",
    "      arr, wight = self.arr, self.wight\n",
    "      ra, rb = self.getRoot(a), self.getRoot(b)\n",
    "      if wight[ra] < wight[rb]:\n",
    "          arr[ra] = rb\n",
    "          wight[rb] += wight[ra]\n",
    "      else:\n",
    "          arr[rb] = ra\n",
    "          wight[ra] += wight[rb]\n",
    "  \n",
    "  def find(self, a, b):\n",
    "      return self.getRoot(a) == self.getRoot(b)\n",
    "  \n",
    "  def gruop(self):\n",
    "      from collections import defaultdict\n",
    "      ans = defaultdict(set)\n",
    "      for i in range(len(self.arr)):\n",
    "          ans[self.getRoot(i)].add(i)\n",
    "      return ans\n",
    "\n",
    "class Solution:\n",
    "  def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "      # 获得质数节点集合\n",
    "      primeLis = set(primes[:bisect_right(primes, n)])\n",
    "      # 建图\n",
    "      G = [[] for i in range(n + 1)]\n",
    "      # 并查集必须有节点加权优化,否则超时\n",
    "      uf = UnionFind(n + 1)\n",
    "      # 建图, 维护连通分量\n",
    "      for a, b in edges:\n",
    "          if a not in primeLis or b not in primeLis:\n",
    "              G[a].append(b)\n",
    "              G[b].append(a)\n",
    "          if (a not in primeLis) and (b not in primeLis):\n",
    "              uf.union(a, b)\n",
    "      # 获得所有连通分量的具体分组\n",
    "      group = uf.gruop()\n",
    "      # 统计每个节点所属组的节点数量\n",
    "      cnt = [0] * (n + 1)\n",
    "      for set_ in group.values():\n",
    "          for i in set_:\n",
    "              cnt[i] = len(set_)\n",
    "      \n",
    "      ans = 0\n",
    "      # 枚举质数节点\n",
    "      for p in primeLis:\n",
    "          num = 0\n",
    "          # 与质数节点相连的组中, 任选两个组的路径都可以相互搭配, 单个组的路径也可\n",
    "          # 可以使用乘法原理求出总路径数\n",
    "          for n in G[p]:\n",
    "              ans += num * cnt[n] + cnt[n]\n",
    "              num += cnt[n]\n",
    "      return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "N = 10**5+1\n",
    "shai = [1] * N\n",
    "zhis = []\n",
    "for i in range(2,N):\n",
    "    if shai[i]:\n",
    "        zhis.append(i)\n",
    "        j = i*i\n",
    "        while j<N:\n",
    "            shai[j] = 0\n",
    "            j += i\n",
    "zhiset = set(zhis)\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        neis = [[] for _ in range(n)]\n",
    "        debug = 0\n",
    "        for x,y in edges:\n",
    "            neis[x-1].append(y-1)\n",
    "            neis[y-1].append(x-1)\n",
    "        visited = [0] * n\n",
    "        q = [0]\n",
    "        levels = []\n",
    "        childs = [[] for _ in range(n)]\n",
    "        while q:\n",
    "            levels.append([])\n",
    "            temp = []\n",
    "            while q:\n",
    "                now = q.pop()\n",
    "                if visited[now]:\n",
    "                    continue\n",
    "                levels[-1].append(now)\n",
    "                visited[now] = 1\n",
    "                for nei in neis[now]:\n",
    "                    if visited[nei]:\n",
    "                        continue\n",
    "                    temp.append(nei)\n",
    "                    childs[now].append(nei)\n",
    "            q = temp\n",
    "        rec = [[0,0] for _ in range(n)]\n",
    "        out = 0\n",
    "        if debug:\n",
    "            print(levels,childs)\n",
    "        while levels:\n",
    "            level = levels.pop()\n",
    "            for now in level:\n",
    "                temp = [0,0]\n",
    "                for child in childs[now]:\n",
    "                    if now+1 in zhiset:\n",
    "                        out += temp[0]*rec[child][0]\n",
    "                    else:\n",
    "                        out += temp[0]*rec[child][1]+temp[1]*rec[child][0]\n",
    "                    for i in range(2):\n",
    "                        temp[i] += rec[child][i]\n",
    "                if now+1 in zhiset:\n",
    "                    out += temp[0]\n",
    "                else:\n",
    "                    out += temp[1]\n",
    "                if now+1 in zhiset:\n",
    "                    rec[now][1] = temp[0]+1\n",
    "                else:\n",
    "                    rec[now] = temp\n",
    "                    rec[now][0] += 1\n",
    "                if debug:\n",
    "                    print(now,out)\n",
    "        if debug:\n",
    "            print(rec)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        tree = [[] for _ in range(n)]\n",
    "        for u, v in edges:\n",
    "            u, v = u - 1, v - 1\n",
    "            tree[u].append((v, 1))\n",
    "            tree[v].append((u, 1))\n",
    "\n",
    "        centTree, root = centroidDecomposition(n, tree)\n",
    "        removed = [False] * n\n",
    "        res = 0\n",
    "\n",
    "        def collect(cur: int, pre: int, sub: \"DefaultDict[int, int]\", primeCount: int) -> None:\n",
    "            \"\"\"统计子树内的答案.\"\"\"\n",
    "            primeCount += int(E.isPrime(cur + 1))\n",
    "            if primeCount <= 1:\n",
    "                sub[primeCount] += 1\n",
    "            else:\n",
    "                return\n",
    "            for next, _ in tree[cur]:\n",
    "                if next != pre and not removed[next]:\n",
    "                    collect(next, cur, sub, primeCount)\n",
    "\n",
    "        def decomposition(cur: int, pre: int) -> None:\n",
    "            nonlocal res\n",
    "            removed[cur] = True\n",
    "            for next in centTree[cur]:  # 点分树的子树中的答案(不经过重心)\n",
    "                if not removed[next]:\n",
    "                    decomposition(next, cur)\n",
    "            removed[cur] = False\n",
    "\n",
    "            init = int(E.isPrime(cur + 1))\n",
    "            counter = defaultdict(int, {init: 1})  # 经过重心的路径\n",
    "            for next, _ in tree[cur]:\n",
    "                if next == pre or removed[next]:\n",
    "                    continue\n",
    "                sub = defaultdict(int)  # 统计子树内0和1的路径数(不包含当前根节点cur)\n",
    "                collect(next, cur, sub, init)\n",
    "                if init == 0:\n",
    "                    res += sub[0] * counter[1]\n",
    "                    res += sub[1] * counter[0]\n",
    "                    counter[0] += sub[0]\n",
    "                    counter[1] += sub[1]\n",
    "                else:\n",
    "                    res += sub[1] * counter[1]\n",
    "                    counter[1] += sub[1]\n",
    "\n",
    "        decomposition(root, -1)\n",
    "        return res\n",
    "\n",
    "\n",
    "\n",
    "  \n",
    "from collections import defaultdict\n",
    "from typing import DefaultDict, List, Tuple\n",
    "from collections import Counter\n",
    "\n",
    "class EratosthenesSieve:\n",
    "    \"\"\"埃氏筛\"\"\"\n",
    "\n",
    "    __slots__ = \"_minPrime\"  # 每个数的最小质因数\n",
    "\n",
    "    def __init__(self, maxN: int):\n",
    "        \"\"\"预处理 O(nloglogn)\"\"\"\n",
    "        minPrime = list(range(maxN + 1))\n",
    "        upper = int(maxN**0.5) + 1\n",
    "        for i in range(2, upper):\n",
    "            if minPrime[i] < i:\n",
    "                continue\n",
    "            for j in range(i * i, maxN + 1, i):\n",
    "                if minPrime[j] == j:\n",
    "                    minPrime[j] = i\n",
    "        self._minPrime = minPrime\n",
    "\n",
    "    def isPrime(self, n: int) -> bool:\n",
    "        if n < 2:\n",
    "            return False\n",
    "        return self._minPrime[n] == n\n",
    "\n",
    "    def getPrimeFactors(self, n: int) -> \"Counter[int]\":\n",
    "        \"\"\"求n的质因数分解 O(logn)\"\"\"\n",
    "        res, f = Counter(), self._minPrime\n",
    "        while n > 1:\n",
    "            m = f[n]\n",
    "            res[m] += 1\n",
    "            n //= m\n",
    "        return res\n",
    "\n",
    "    def getPrimes(self) -> List[int]:\n",
    "        return [x for i, x in enumerate(self._minPrime) if i >= 2 and i == x]\n",
    "\n",
    "\n",
    "E = EratosthenesSieve(int(1e5) + 10)\n",
    "\n",
    "\n",
    "def centroidDecomposition(n: int, tree: List[List[Tuple[int, int]]]) -> Tuple[List[List[int]], int]:\n",
    "    \"\"\"\n",
    "    树的重心分解, 返回点分树和点分树的根.\n",
    "\n",
    "    Params:\n",
    "        n: 树的节点数.\n",
    "        tree: `无向树`的邻接表.\n",
    "    Returns:\n",
    "        centTree: 重心互相连接形成的有根树, 可以想象把树拎起来, 重心在树的中心，连接着各个子树的重心...\n",
    "        root: 点分树的根.\n",
    "    \"\"\"\n",
    "\n",
    "    def getSize(cur: int, parent: int) -> int:\n",
    "        subSize[cur] = 1\n",
    "        for next, _ in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            subSize[cur] += getSize(next, cur)\n",
    "        return subSize[cur]\n",
    "\n",
    "    def getCentroid(cur: int, parent: int, mid: int) -> int:\n",
    "        for next, _ in tree[cur]:\n",
    "            if next == parent or removed[next]:\n",
    "                continue\n",
    "            if subSize[next] > mid:\n",
    "                return getCentroid(next, cur, mid)\n",
    "        return cur\n",
    "\n",
    "    def build(cur: int) -> int:\n",
    "        centroid = getCentroid(cur, -1, getSize(cur, -1) // 2)\n",
    "        removed[centroid] = True\n",
    "        for next, _ in tree[centroid]:\n",
    "            if not removed[next]:\n",
    "                centTree[centroid].append(build(next))\n",
    "        removed[centroid] = False\n",
    "        return centroid\n",
    "\n",
    "    subSize = [0] * n\n",
    "    removed = [False] * n\n",
    "    centTree = [[] for _ in range(n)]\n",
    "    root = build(0)\n",
    "    return centTree, root\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self.root = [i for i in range(n)]\n",
    "        self.size = [1] * n\n",
    "        self.part = n\n",
    "        return\n",
    "\n",
    "    def find(self, x):\n",
    "        lst = []\n",
    "        while x != self.root[x]:\n",
    "            lst.append(x)\n",
    "            # 在查询的时候合并到顺带直接根节点\n",
    "            x = self.root[x]\n",
    "        for w in lst:\n",
    "            self.root[w] = x\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        root_x = self.find(x)\n",
    "        root_y = self.find(y)\n",
    "        if root_x == root_y:\n",
    "            return False\n",
    "        if self.size[root_x] >= self.size[root_y]:\n",
    "            root_x, root_y = root_y, root_x\n",
    "        self.root[root_x] = root_y\n",
    "        self.size[root_y] += self.size[root_x]\n",
    "        # 将非根节点的秩赋0\n",
    "        self.size[root_x] = 0\n",
    "        self.part -= 1\n",
    "        return True\n",
    "\n",
    "    def is_connected(self, x, y):\n",
    "        return self.find(x) == self.find(y)\n",
    "\n",
    "    def get_root_part(self):\n",
    "        # 获取每个根节点对应的组\n",
    "        part = defaultdict(list)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            part[self.find(i)].append(i)\n",
    "        return part\n",
    "\n",
    "    def get_root_size(self):\n",
    "        # 获取每个根节点对应的组大小\n",
    "        size = defaultdict(int)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            size[self.find(i)] = self.size[self.find(i)]\n",
    "        return size\n",
    "\n",
    "\n",
    "class NumberTheoryPrimeFactor:\n",
    "    def __init__(self, ceil):\n",
    "        self.ceil = ceil + 100\n",
    "        self.prime_factor = [[] for _ in range(self.ceil + 1)]\n",
    "        self.min_prime = [0] * (self.ceil + 1)\n",
    "        self.is_prime = [0] * (self.ceil + 1)\n",
    "        self.get_min_prime_and_prime_factor()\n",
    "        return\n",
    "\n",
    "    def get_min_prime_and_prime_factor(self):\n",
    "        # 模板：计算 1 到 self.ceil 所有数字的最小质数因子\n",
    "        for i in range(2, self.ceil + 1):\n",
    "            if not self.min_prime[i]:\n",
    "                self.is_prime[i] = 1\n",
    "                self.min_prime[i] = i\n",
    "                for j in range(i * i, self.ceil + 1, i):\n",
    "                    if not self.min_prime[j]:\n",
    "                        self.min_prime[j] = i\n",
    "\n",
    "        # 模板：计算 1 到 self.ceil 所有数字的质数分解（可选）\n",
    "        for num in range(2, self.ceil + 1):\n",
    "            i = num\n",
    "            while num > 1:\n",
    "                p = self.min_prime[num]\n",
    "                cnt = 0\n",
    "                while num % p == 0:\n",
    "                    num //= p\n",
    "                    cnt += 1\n",
    "                self.prime_factor[i].append([p, cnt])\n",
    "        return\n",
    "    \n",
    "nt = NumberTheoryPrimeFactor(10**5)\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        dct = [[] for _ in range(n+1)]\n",
    "        for i, j in edges:\n",
    "            dct[i].append(j)\n",
    "            dct[j].append(i)\n",
    "        stack = [[1, -1]]\n",
    "        ans = 0\n",
    "        sub0 = [0]*(n+1)\n",
    "        sub1 = [0]*(n+1)\n",
    "        while stack:\n",
    "            i, fa = stack.pop()\n",
    "            if i >= 0:\n",
    "                stack.append([~i, fa])\n",
    "                for j in dct[i]:\n",
    "                    if j != fa:\n",
    "                        stack.append([j, i])\n",
    "            else:\n",
    "                i = ~i\n",
    "                cnt0 = cnt1 = 0\n",
    "                for j in dct[i]:\n",
    "                    if j != fa:\n",
    "                        if nt.is_prime[i]:\n",
    "                            ans += cnt0*sub0[j]\n",
    "                        else:\n",
    "                            ans += cnt0*sub1[j] + cnt1*sub0[j]\n",
    "                        cnt0 += sub0[j]\n",
    "                        cnt1 += sub1[j]\n",
    "                        \n",
    "                if nt.is_prime[i]:\n",
    "                    ans += cnt0\n",
    "                    sub0[i] = 0\n",
    "                    sub1[i] = cnt0+1\n",
    "                else:\n",
    "                    ans += cnt1\n",
    "                    sub0[i] = cnt0+1\n",
    "                    sub1[i] = cnt1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        ee = defaultdict(list)\n",
    "\n",
    "        for x, y in edges:\n",
    "            ee[x].append(y)\n",
    "            ee[y].append(x)\n",
    "\n",
    "        primes = [True] * (n + 1)\n",
    "        primes[1] = primes[0] = False\n",
    "        for i in range(2, n):\n",
    "            j = i\n",
    "            while i * j <= n:\n",
    "                primes[j * i] = False\n",
    "                j += 1\n",
    "\n",
    "        @cache    \n",
    "        def dfs(cur, pa):\n",
    "            if primes[cur]:\n",
    "                return 0\n",
    "            ans = 1\n",
    "            for nei in ee[cur]:\n",
    "                if nei == pa:\n",
    "                    continue\n",
    "                ans += dfs(nei, cur)\n",
    "            # print(\"size[\", cur, \"] = \", ans)\n",
    "            return ans\n",
    "        \n",
    "        curr = 0\n",
    "        # print(primes)\n",
    "        # print(size)\n",
    "        # print(ee)\n",
    "        for i in range(1, n+1):\n",
    "            if primes[i]:\n",
    "                s = 0\n",
    "                for nei in ee[i]:\n",
    "                    nn = dfs(nei, -1)\n",
    "                    curr += s * nn\n",
    "                    s += nn\n",
    "                curr += s\n",
    "        return curr\n",
    "\n",
    "                \n",
    "                \n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Prime:\n",
    "    def prime_sieve(self, n):\n",
    "        \"\"\"returns a sieve of primes >= 5 and < n\"\"\"\n",
    "        flag = n % 6 == 2\n",
    "        sieve = bytearray((n // 3 + flag >> 3) + 1)\n",
    "        for i in range(1, int(n ** 0.5) // 3 + 1):\n",
    "            if not (sieve[i >> 3] >> (i & 7)) & 1:\n",
    "                k = (3 * i + 1) | 1\n",
    "                for j in range(k * k // 3, n // 3 + flag, 2 * k):\n",
    "                    sieve[j >> 3] |= 1 << (j & 7)\n",
    "                for j in range(k * (k - 2 * (i & 1) + 4) // 3, n // 3 + flag, 2 * k):\n",
    "                    sieve[j >> 3] |= 1 << (j & 7)\n",
    "        return sieve\n",
    "\n",
    "    def prime_list(self, n):\n",
    "        \"\"\"returns a list of primes <= n\"\"\"\n",
    "        res = []\n",
    "        if n > 1:\n",
    "            res.append(2)\n",
    "        if n > 2:\n",
    "            res.append(3)\n",
    "        if n > 4:\n",
    "            sieve = self.prime_sieve(n + 1)\n",
    "            res.extend(\n",
    "                3 * i + 1 | 1 for i in range(1, (n + 1) // 3 + (n % 6 == 1)) if not (sieve[i >> 3] >> (i & 7)) & 1)\n",
    "        return res\n",
    "\n",
    "    def __init__(self, n) -> None:\n",
    "        self.primes = self.prime_list(n)\n",
    "\n",
    "    def dissolve(self, num):\n",
    "        '''prime factor decomposition of num'''\n",
    "        lst = []\n",
    "        idx = -1\n",
    "        for prime in self.primes:\n",
    "            if prime * prime > num:\n",
    "                break\n",
    "\n",
    "            if num % prime == 0:\n",
    "                lst.append([prime, 0])\n",
    "                idx += 1\n",
    "\n",
    "            while num % prime == 0:\n",
    "                lst[idx][1] += 1\n",
    "                num //= prime\n",
    "\n",
    "        if num != 1:\n",
    "            lst.append([num, 1])\n",
    "\n",
    "        return lst\n",
    "\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        ps = set(Prime(10 ** 5 + 5).primes)\n",
    "        g = [[] for _ in range(n + 1)]\n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        @cache\n",
    "        def dfs(x, p):\n",
    "            if x in ps:\n",
    "                return 0\n",
    "            ans = 1\n",
    "            for y in g[x]:\n",
    "                if y == p:\n",
    "                    continue\n",
    "                ans += dfs(y, x)\n",
    "            return ans\n",
    "\n",
    "        res = 0\n",
    "        for x in range(1, n + 1):\n",
    "            if x in ps:\n",
    "                s = 0\n",
    "                for y in g[x]:\n",
    "                    tmp = dfs(y, x)\n",
    "                    res += tmp * s\n",
    "                    s += tmp\n",
    "                res += s\n",
    "                \n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import defaultdict\n",
    "\n",
    "def sieve_of_sundaram(mx):\n",
    "    if mx < 2:\n",
    "        return []\n",
    "    mx_idx = (mx - 1) // 2\n",
    "    is_prime_idx = [True] * (mx_idx + 1)\n",
    "    for i in range(1, len(is_prime_idx)):\n",
    "        if 2 * i * i + i + i > mx_idx:\n",
    "            break\n",
    "        for j in range(i, len(is_prime_idx)):\n",
    "            if 2 * i * j + i + j > mx_idx:\n",
    "                break\n",
    "            else:\n",
    "                is_prime_idx[2 * i * j + i + j] = False\n",
    "    return [2] + [2 * i + 1 for i, is_prime in enumerate(is_prime_idx) if is_prime and i > 0]\n",
    "\n",
    "PRIMES = sieve_of_sundaram(10 ** 5)\n",
    "PRIMES_SET = set(PRIMES)\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        def build_graph(edges):\n",
    "            graph = defaultdict(list)\n",
    "            for u, v in edges:\n",
    "                graph[u].append(v)\n",
    "                graph[v].append(u)\n",
    "            return graph\n",
    "\n",
    "        @cache\n",
    "        def dfs(node, parent):\n",
    "            if node in PRIMES_SET:\n",
    "                return 0\n",
    "            res = 1\n",
    "            for nxt in graph[node]:\n",
    "                if nxt != parent:\n",
    "                    res += dfs(nxt, node)\n",
    "            return res\n",
    "        graph = build_graph(edges)\n",
    "        res = 0\n",
    "        for prime in PRIMES:\n",
    "            if prime > n:\n",
    "                break\n",
    "            cnt_of_a = 1\n",
    "            cnt_of_path = 0\n",
    "            for neighbor in graph[prime]:\n",
    "                branch = dfs(neighbor, prime)\n",
    "                cnt_of_path += cnt_of_a * branch\n",
    "                cnt_of_a += branch\n",
    "            res += cnt_of_path\n",
    "        return res\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "\n",
    "        if n == 1: return 0\n",
    "\n",
    "        p = []\n",
    "        \n",
    "        def isprime(x):\n",
    "            for v in p:\n",
    "                if v * v > x: return True\n",
    "                if x % v == 0: return False\n",
    "            return True\n",
    "        \n",
    "        for i in range(2, n+1):\n",
    "            if isprime(i): p.append(i)\n",
    "        \n",
    "        g = defaultdict(list)\n",
    "        \n",
    "        for a, b in edges:\n",
    "            g[a].append(b)\n",
    "            g[b].append(a)\n",
    "        \n",
    "        p = set(p)\n",
    "        k = defaultdict(dict)\n",
    "        s = [0] * (n+1)\n",
    "        def dfs(x, f):\n",
    "            cnt = 0\n",
    "            for y in g[x]:\n",
    "                if y == f: continue\n",
    "                r = dfs(y, x)\n",
    "                k[x][y] = r\n",
    "                cnt += r\n",
    "            s[x] = cnt\n",
    "            if x in p: return 0\n",
    "            return cnt + 1\n",
    "        \n",
    "        def update(x, f):\n",
    "            if f > 0:\n",
    "                if f in p: k[x][f] = 0\n",
    "                else:\n",
    "                    if x in p:\n",
    "                        k[x][f] = s[f] + 1\n",
    "                        s[x] += s[f] + 1\n",
    "                    else:\n",
    "                        k[x][f] = s[f] - s[x]\n",
    "                        s[x] = s[f]\n",
    "            for y in g[x]:\n",
    "                if y == f: continue\n",
    "                update(y, x)\n",
    "        dfs(2, -1)\n",
    "        update(2, -1)\n",
    "        #print(s)\n",
    "        #print(k)\n",
    "        ans = 0\n",
    "        for v in p:\n",
    "            ans += s[v]\n",
    "            cnt = 0\n",
    "            for x in k[v]:\n",
    "                cnt += k[v][x] * (s[v] - k[v][x])\n",
    "            ans += cnt // 2\n",
    "        return ans\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from math import gcd\n",
    "import bisect\n",
    "import copy\n",
    "import decimal\n",
    "import fractions\n",
    "import heapq\n",
    "import itertools\n",
    "import math\n",
    "import random\n",
    "import sys\n",
    "import time\n",
    "from collections import Counter, deque, defaultdict\n",
    "from functools import lru_cache, reduce\n",
    "from heapq import heappush, heappop, heapify, heappushpop\n",
    "\n",
    "\n",
    "def I(): return map(int, input().split())\n",
    "\n",
    "from typing import List\n",
    "\n",
    "class PrimeTable:\n",
    "    def __init__(self, n:int) -> None:\n",
    "        self.n = n\n",
    "        self.primes = []\n",
    "        self.max_div = list(range(n+1))\n",
    "        self.max_div[1] = 1\n",
    "        self.phi = list(range(n+1))\n",
    "\n",
    "        for i in range(2, n + 1):\n",
    "            if self.max_div[i] == i:\n",
    "                self.primes.append(i)\n",
    "                for j in range(i, n+1, i):\n",
    "                    self.max_div[j] = i\n",
    "                    self.phi[j] = self.phi[j] // i * (i-1)\n",
    "\n",
    "    def is_prime(self, x:int):\n",
    "        if x < 2: return False\n",
    "        if x <= self.n: return self.max_div[x] == x\n",
    "        for p in self.primes:\n",
    "            if p * p > x: break\n",
    "            if x % p == 0: return False\n",
    "        return True\n",
    "\n",
    "    def prime_factorization(self, x:int):\n",
    "        if x > self.n:\n",
    "            for p in self.primes:\n",
    "                if p * p > x: break\n",
    "                if x <= self.n: break\n",
    "                if x % p == 0:\n",
    "                    cnt = 0\n",
    "                    while x % p == 0: cnt += 1; x //= p\n",
    "                    yield p, cnt\n",
    "        while (1 < x and x <= self.n):\n",
    "            p, cnt = self.max_div[x], 0\n",
    "            while x % p == 0: cnt += 1; x //= p\n",
    "            yield p, cnt\n",
    "        if x >= self.n and x > 1:\n",
    "            yield x, 1\n",
    "\n",
    "    def get_factors(self, x:int):\n",
    "        factors = [1]\n",
    "        for p, b in self.prime_factorization(x):\n",
    "            n = len(factors)\n",
    "            for j in range(1, b+1):\n",
    "                for d in factors[:n]:\n",
    "                    factors.append(d * (p ** j))\n",
    "        return factors\n",
    "\n",
    "pt = PrimeTable(10 ** 6 + 1)\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            g[u].append(v)\n",
    "            g[v].append(u)\n",
    "        \n",
    "        def dfs(u, p, nodes):\n",
    "            nodes.append(u)\n",
    "            for v in g[u]:\n",
    "                if v == p: continue\n",
    "                if not pt.is_prime(v):\n",
    "                    dfs(v, u, nodes)\n",
    "        ans = 0\n",
    "        sz = [0] * (n + 1)\n",
    "        for x in range(1, n + 1):\n",
    "            if not pt.is_prime(x): continue\n",
    "            else:\n",
    "                curr = 0\n",
    "                for y in g[x]:\n",
    "                    if pt.is_prime(y): continue\n",
    "                    if sz[y] == 0:\n",
    "                        nodes = []\n",
    "                        dfs(y, -1, nodes)\n",
    "                        for node in nodes:\n",
    "                            sz[node] = len(nodes)\n",
    "                    ans += sz[y] * curr\n",
    "                    curr += sz[y]\n",
    "                ans += curr\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "N = 10**5\n",
    "is_prime = [True] * (N+1)\n",
    "is_prime[0] = is_prime[1] = False\n",
    "\n",
    "for i in range(2, int(N**0.5) + 1):\n",
    "    if is_prime[i]:\n",
    "        for j in range(i*i, N+1, i):\n",
    "            is_prime[j] = False\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        if n == 1:\n",
    "            return 0\n",
    "        \n",
    "        neb = defaultdict(list)\n",
    "        for u, v in edges:\n",
    "            neb[u].append(v)\n",
    "            neb[v].append(u)\n",
    "        \n",
    "        def dfs(u, fa):\n",
    "            cnt = 1\n",
    "            for v in neb[u]:\n",
    "                if v != fa:\n",
    "                    r = dfs(v, u)\n",
    "                    if not is_prime[v]:\n",
    "                        cnt += r\n",
    "            gg[u] = cnt\n",
    "            return cnt\n",
    "        gg = [0] * (n+1)\n",
    "        dfs(2, -1)\n",
    "        #print(f'{gg=}')\n",
    "        \n",
    "        \n",
    "        def dfs2(u, fa):\n",
    "            old_u = gg[u]\n",
    "            #print(f'now {u=} {gg=}')\n",
    "            if is_prime[u]:\n",
    "                if len(neb[u]) == 1:\n",
    "                    cnt = gg[neb[u][0]] if not is_prime[neb[u][0]] else 0\n",
    "                    #print(f'for {u=} {cnt=} {gg=}')\n",
    "                else:\n",
    "                    v = neb[u][0]\n",
    "                    prev = gg[v] if not is_prime[v] else 0\n",
    "                    cnt = 0\n",
    "                    valid = prev > 0 \n",
    "                    #print(f'for {u=} {gg=}')\n",
    "                    for v in neb[u][1:]:\n",
    "                        if not is_prime[v]:\n",
    "                            #print(f'{prev=} {gg[v]=}')\n",
    "                            cnt += (prev) * (gg[v])\n",
    "                            prev += gg[v]\n",
    "                            valid = True\n",
    "                    cnt += old_u-1\n",
    "                    #print(f'for {u=} add {cnt}')\n",
    "                self._res += cnt\n",
    "            \n",
    "            for v in neb[u]:\n",
    "                if v != fa:\n",
    "                        \n",
    "                    old_v = gg[v]\n",
    "                    \n",
    "                    gg[u] -= old_v if not is_prime[v] else 0\n",
    "                    gg[v] += gg[u] if not is_prime[u] else 0\n",
    "                    \n",
    "                    dfs2(v, u)\n",
    "                    \n",
    "                    gg[u] = old_u\n",
    "                    gg[v] = old_v\n",
    "            \n",
    "        self._res = 0\n",
    "        dfs2(2, -1)\n",
    "        \n",
    "        return self._res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def prime_gen(N):\n",
    "    bool_list = [False] * (N + 1)\n",
    "    bool_list[0] = bool_list[1] = True\n",
    "    p_list = []\n",
    "    for i in range(2, N+1):\n",
    "        if not bool_list[i]:\n",
    "            p_list.append(i)\n",
    "        for p in p_list:\n",
    "            j = p * i\n",
    "            if j > N:\n",
    "                break\n",
    "            bool_list[j] = True\n",
    "            if i % p == 0:\n",
    "                break\n",
    "    ret_list = [False] * N\n",
    "    for p in p_list:\n",
    "        ret_list[p-1] = True\n",
    "    return ret_list \n",
    "prime_list = prime_gen(100000)\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parent = list(range(n))\n",
    "    def find(self, x):\n",
    "        if self.parent[x] != x:\n",
    "            self.parent[x] = self.find(self.parent[x])\n",
    "        return self.parent[x]\n",
    "    def union(self, x, y):\n",
    "        p_x = self.find(x)\n",
    "        p_y = self.find(y)\n",
    "        if p_x != p_y:\n",
    "            self.parent[p_x] = p_y\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        map_dict = collections.defaultdict(list)\n",
    "        uf_obj = UnionFind(n)\n",
    "        for i, j in edges:\n",
    "            i -= 1\n",
    "            j -= 1\n",
    "            if prime_list[i] or prime_list[j]:\n",
    "                if prime_list[i]:\n",
    "                    map_dict[i].append(j)\n",
    "                if prime_list[j]:\n",
    "                    map_dict[j].append(i)\n",
    "            else:\n",
    "                uf_obj.union(i, j)\n",
    "        count_dict = collections.defaultdict(int)\n",
    "        for i in range(n):\n",
    "            if not prime_list[i]:\n",
    "                count_dict[uf_obj.find(i)] += 1\n",
    "        #print(map_dict)\n",
    "        #print(count_dict)\n",
    "        ret = 0\n",
    "        for val_list in map_dict.values():\n",
    "            tmp_list = [count_dict[uf_obj.find(j)] for j in val_list if not prime_list[j]]\n",
    "            if not tmp_list:\n",
    "                continue\n",
    "            prev = tmp_list[0]\n",
    "            for i in range(1, len(tmp_list)):\n",
    "                ret += tmp_list[i] * prev\n",
    "                prev += tmp_list[i]\n",
    "            ret += prev\n",
    "        return ret\n",
    "\n",
    "\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        if n < 3:\n",
    "            return n - 1\n",
    "        a = [3]\n",
    "        for i in range(5, n + 1, 2):\n",
    "            t = isqrt(i)\n",
    "            for j in a:\n",
    "                if j > t:\n",
    "                    a.append(i)\n",
    "                    break\n",
    "                if i % j == 0:\n",
    "                    break\n",
    "        s = set(a)\n",
    "        s.add(2)\n",
    "        a = list(range(n + 1))\n",
    "        for i in s:\n",
    "            a[i] = []\n",
    "            \n",
    "        def f(i):\n",
    "            if a[i] != i:\n",
    "                a[i] = f(a[i])\n",
    "            return a[i]\n",
    "        \n",
    "        for i, j in edges:\n",
    "            if i in s:\n",
    "                if j not in s:\n",
    "                    a[i].append(j)\n",
    "            elif j in s:\n",
    "                a[j].append(i)\n",
    "            else:\n",
    "                a[f(i)] = f(j)\n",
    "        d = {i: f(i) for i in range(1, n + 1) if i not in s}\n",
    "        c = Counter(d.values())\n",
    "        for i in d:\n",
    "            a[i] = c[d[i]]\n",
    "        ans = 0\n",
    "        for i in s:\n",
    "            c = Counter(a[j] for j in a[i])\n",
    "            ans += sum(i * i * comb(j, 2) + i * j for i, j in c.items()) + sum(i * j * c[i] * c[j] for i, j in combinations(c, 2))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        if n < 3:\n",
    "            return n - 1\n",
    "            \n",
    "        def f():\n",
    "            yield 2\n",
    "            m = n - 1 >> 1\n",
    "            a = [True] * m\n",
    "            t = int((sqrt(2 * m + 1) - 3) / 2)\n",
    "            for i in range(m):\n",
    "                if a[i]:\n",
    "                    if i <= t:\n",
    "                        for j in range(i * i * 2 + 6 * i + 3, m, 2 * i + 3):\n",
    "                            a[j] = False\n",
    "                    yield 2 * i + 3\n",
    "            \n",
    "        s = set(f())\n",
    "        a = list(range(n + 1))\n",
    "        for i in s:\n",
    "            a[i] = []\n",
    "            \n",
    "        def f(i):\n",
    "            if a[i] != i:\n",
    "                a[i] = f(a[i])\n",
    "            return a[i]\n",
    "        \n",
    "        for i, j in edges:\n",
    "            if i in s:\n",
    "                if j not in s:\n",
    "                    a[i].append(j)\n",
    "            elif j in s:\n",
    "                a[j].append(i)\n",
    "            else:\n",
    "                a[f(i)] = f(j)\n",
    "        d = Counter(f(i) for i in range(1, n + 1) if i not in s)       \n",
    "        ans = 0\n",
    "        for i in s:\n",
    "            c = Counter(d[a[j]] for j in a[i])\n",
    "            ans += sum(i * i * comb(j, 2) + i * j for i, j in c.items()) + sum(i * j * c[i] * c[j] for i, j in combinations(c, 2))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "N = 10 ** 5 + 5\n",
    "minp = [0] * N\n",
    "primes = []\n",
    "for i in range(2, N):\n",
    "    if minp[i] == 0:\n",
    "        minp[i] = i\n",
    "        primes.append(i)\n",
    "    for p in primes:\n",
    "        if p * i >= N: break\n",
    "        minp[i * p] = p\n",
    "        if minp[i] == p: break\n",
    "            \n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int:\n",
    "        g = [[] for _ in range(n + 1)]\n",
    "        for x, y in edges:\n",
    "            g[x].append(y)\n",
    "            g[y].append(x)\n",
    "        Siz = [0] * (n + 1)\n",
    "        def dfs(root, Fa):\n",
    "            Siz[root] = 1\n",
    "            for y in g[root]:\n",
    "                if y != Fa:\n",
    "                    Siz[root] += dfs(y, root)\n",
    "            if minp[root] == root: return 0\n",
    "            return Siz[root]\n",
    "        dfs(1, 0)\n",
    "        def DFS(root, fa, cnt):\n",
    "            ans = 0\n",
    "            if minp[root] == root:\n",
    "                x = cnt + Siz[root]\n",
    "                ans += x * (x - 1) // 2 - cnt * (cnt - 1) // 2\n",
    "                for y in g[root]:\n",
    "                    if y != fa:\n",
    "                        if minp[y] != y:\n",
    "                            ans -= Siz[y] * (Siz[y] - 1) // 2\n",
    "                            ans += DFS(y, root, Siz[y])\n",
    "                        else:\n",
    "                            ans += DFS(y, root, 0)\n",
    "            else:\n",
    "                for y in g[root]:\n",
    "                    if y != fa:\n",
    "                        ans += DFS(y, root, cnt)\n",
    "            return ans\n",
    "        return DFS(1, 0, Siz[1])\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "N = 10**5 + 1\n",
    "prime = [True] * N\n",
    "prime[0] = prime[1] = False\n",
    "for i in range(2, N):\n",
    "    if prime[i]:\n",
    "        for j in range(i*i, N, i):\n",
    "            prime[j] = False\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def countPaths(self, n: int, edges: List[List[int]]) -> int: \n",
    "        fa = list(range(n+1))\n",
    "        def find(i):\n",
    "            if fa[i] != i:\n",
    "                fa[i] = find(fa[i])\n",
    "            return fa[i]\n",
    "        \n",
    "        G = [[] for _ in range(n+1)]\n",
    "        for u,v in edges:\n",
    "            G[u].append(v)\n",
    "            G[v].append(u)\n",
    "            if prime[u] == prime[v] == False:\n",
    "                fu, fv = find(u), find(v)\n",
    "                fa[fu] = fa[fv]\n",
    "        for i in range(n+1):\n",
    "            find(i)\n",
    "        c = Counter(fa)\n",
    "\n",
    "        ans = 0\n",
    "        def dfs(i, p):\n",
    "            nonlocal ans\n",
    "            s = 0\n",
    "            for j in G[i]:\n",
    "                if prime[i] and not prime[j]:\n",
    "                    ans += c[fa[j]] * (s+1)\n",
    "                    s += c[fa[j]]\n",
    "                if j != p:\n",
    "                    dfs(j, i)\n",
    "\n",
    "        dfs(1, -1)\n",
    "        return ans\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
